#!/usr/bin/env node
"use strict";

const { spawn } = require("child_process");
const fs = require("fs");
const path = require("path");

function fileExists(p) {
  try {
    fs.accessSync(p, fs.constants.F_OK);
    return true;
  } catch {
    return false;
  }
}

function resolveWhisperExe() {
  const localExe = path.join("venv", "Scripts", "whisper.exe");
  if (fileExists(localExe)) return localExe;
  // fallback: try global whisper in PATH
  return "whisper";
}

function toWinPath(p) {
  return p.replace(/\\/g, "/");
}

async function main() {
  const inputArgIdx = process.argv.indexOf("-i");
  const inputPath = inputArgIdx > -1 ? process.argv[inputArgIdx + 1] : "assets/lihun.mp4";
  const modelArgIdx = process.argv.indexOf("-m");
  const modelName = modelArgIdx > -1 ? process.argv[modelArgIdx + 1] : "medium";
  const outDirIdx = process.argv.indexOf("-o");
  const outDir = outDirIdx > -1 ? process.argv[outDirIdx + 1] : path.dirname(inputPath) || ".";

  if (!fileExists(inputPath)) {
    console.error(`❌ 输入文件不存在: ${inputPath}`);
    process.exit(1);
  }

  const absInput = path.resolve(inputPath);
  const absOutDir = path.resolve(outDir);
  if (!fs.existsSync(absOutDir)) fs.mkdirSync(absOutDir, { recursive: true });

  const whisperExe = resolveWhisperExe();
  const args = [
    toWinPath(absInput),
    "--model", modelName,
    "--language", "zh",
    "--task", "transcribe",
    "--output_format", "srt",
    "--output_dir", toWinPath(absOutDir),
    "--condition_on_previous_text", "True",
  ];

  console.log(`▶️  执行: ${whisperExe} ${args.join(" ")}`);

  const child = spawn(whisperExe, args, { stdio: ["ignore", "pipe", "pipe"], shell: process.platform === "win32" });

  child.stdout.on("data", (d) => process.stdout.write(d));
  child.stderr.on("data", (d) => process.stderr.write(d));

  child.on("close", (code) => {
    if (code !== 0) {
      console.error(`❌ whisper 进程退出码: ${code}`);
      process.exit(code);
    }

    const base = path.basename(absInput, path.extname(absInput));
    const srtPath = path.join(absOutDir, `${base}.srt`);
    if (fileExists(srtPath)) {
      try {
        // 读取并进行标点感知断句 + 按最大时长拆分片段
        const maxChars = Number(process.env.MAX_LINE_WIDTH || 14);
        const maxLinesPerBlock = Number(process.env.MAX_LINE_COUNT || 2);
        const maxSegmentSecs = Number(process.env.MAX_SEGMENT_SECS || 3);
        const raw = fs.readFileSync(srtPath, "utf8");

        const lines = raw.split(/\r?\n/);
        const blocks = [];
        let i = 0;
        while (i < lines.length) {
          // index line
          let idxLine = "";
          while (i < lines.length && !lines[i].trim()) i++;
          if (i >= lines.length) break;
          idxLine = lines[i++].trim();
          if (!idxLine) continue;

          const timeLine = (i < lines.length ? lines[i++] : "").trim();
          const textLines = [];
          while (i < lines.length && lines[i].trim() !== "") {
            textLines.push(lines[i++]);
          }
          // skip blank separator
          if (i < lines.length && lines[i].trim() === "") i++;

          const text = textLines.join("").replace(/\s+/g, "").trim();
          blocks.push({ time: timeLine, text });
        }

        const splitByPunct = (text) => {
          const parts = text.split(/([，。！？；：,.!?、])/);
          const phrases = [];
          for (let j = 0; j < parts.length; j += 2) {
            const seg = (parts[j] || "").trim();
            const p = parts[j + 1] || "";
            const phrase = (seg + p).trim();
            if (phrase) phrases.push(phrase);
          }
          return phrases;
        };

        const wrapUnlimited = (phrases, limit) => {
          const out = [];
          let current = "";
          const pushCurrent = () => { if (current) { out.push(current); current = ""; } };
          for (const ph of phrases) {
            if (ph.length > limit) {
              pushCurrent();
              for (let k = 0; k < ph.length; k += limit) {
                out.push(ph.slice(k, k + limit));
              }
              continue;
            }
            if (current.length + ph.length <= limit) {
              current += ph;
            } else {
              pushCurrent();
              current = ph;
            }
          }
          pushCurrent();
          return out;
        };

        const rebalanceLines = (arr, limit) => {
          let lines = arr.filter(l => l && l.trim().length > 0);
          // 多次迭代，尽力消除 1 字行
          for (let pass = 0; pass < 3; pass++) {
            for (let i = 0; i < lines.length; i++) {
              const li = lines[i] || "";
              if (li.length !== 1) continue;

              // 1) 优先与前一行合并（不超限）
              if (i > 0 && (lines[i - 1] || "").length + li.length <= limit) {
                lines[i - 1] = (lines[i - 1] || "") + li;
                lines.splice(i, 1); i--; continue;
              }
              // 2) 尝试与后一行合并（不超限）
              if (i + 1 < lines.length && li.length + (lines[i + 1] || "").length <= limit) {
                lines[i] = li + (lines[i + 1] || "");
                lines.splice(i + 1, 1); i--; continue;
              }
              // 3) 从前一行借 1 个字符
              if (i > 0 && (lines[i - 1] || "").length > 2 && li.length + 1 <= limit) {
                const prev = lines[i - 1];
                lines[i - 1] = prev.slice(0, -1);
                lines[i] = prev.slice(-1) + li;
                continue;
              }
              // 4) 从后一行借 1 个字符
              if (i + 1 < lines.length && (lines[i + 1] || "").length > 2 && li.length + 1 <= limit) {
                const next = lines[i + 1];
                lines[i] = li + next.slice(0, 1);
                lines[i + 1] = next.slice(1);
                continue;
              }
              // 5) 兜底：与前一行强制合并（可能略超限，但避免单字行）
              if (i > 0) {
                lines[i - 1] = (lines[i - 1] || "") + li;
                lines.splice(i, 1); i--; continue;
              }
              // 6) 或与后一行强制合并
              if (i + 1 < lines.length) {
                lines[i] = li + (lines[i + 1] || "");
                lines.splice(i + 1, 1); i--; continue;
              }
            }
          }
          return lines;
        };

        const parseTime = (t) => {
          // 00:00:00,000 --> 00:00:03,000
          const [s, e] = t.split(/\s*-->\s*/);
          const toMs = (x) => {
            const m = x.match(/(\d+):(\d+):(\d+),(\d+)/);
            if (!m) return 0;
            return (
              Number(m[1]) * 3600000 +
              Number(m[2]) * 60000 +
              Number(m[3]) * 1000 +
              Number(m[4])
            );
          };
          return { start: toMs(s), end: toMs(e) };
        };

        const fmtTime = (ms) => {
          const pad = (n, w) => String(n).padStart(w, "0");
          const h = Math.floor(ms / 3600000); ms %= 3600000;
          const m = Math.floor(ms / 60000); ms %= 60000;
          const s = Math.floor(ms / 1000); const ms3 = ms % 1000;
          return `${pad(h,2)}:${pad(m,2)}:${pad(s,2)},${pad(ms3,3)}`;
        };

        const newBlocks = [];
        for (const b of blocks) {
          const { start, end } = parseTime(b.time);
          if (!(end > start)) continue;
          const duration = end - start;
          const phrases = splitByPunct(b.text);
          let allLines = wrapUnlimited(phrases, maxChars);
          allLines = rebalanceLines(allLines, maxChars);

          const linesPerSeg = Math.max(1, maxLinesPerBlock);
          const timeSegs = Math.max(1, Math.ceil(duration / (maxSegmentSecs * 1000)));
          const lineSegs = Math.max(1, Math.ceil(allLines.length / linesPerSeg));
          const segCount = Math.max(timeSegs, lineSegs);
          const segDur = Math.max(1, Math.floor(duration / segCount));

          let idx = 0;
          for (let s = 0; s < segCount; s++) {
            const segStart = start + Math.floor((duration * s) / segCount);
            const segEnd = s === segCount - 1 ? end : start + Math.floor((duration * (s + 1)) / segCount);

            const remainingSegs = segCount - s;
            const remainingLines = allLines.length - idx;
            let take = Math.min(linesPerSeg, Math.ceil(remainingLines / remainingSegs));
            if (take < 0) take = 0;
            const linesThis = allLines.slice(idx, idx + take);
            idx += take;

            if (linesThis.length === 0) continue; // 跳过空内容片段
            newBlocks.push({ start: segStart, end: segEnd, lines: linesThis });
          }
        }

        // 重置索引并写回
        const out = [];
        for (let k = 0; k < newBlocks.length; k++) {
          const nb = newBlocks[k];
          out.push(String(k + 1));
          out.push(`${fmtTime(nb.start)} --> ${fmtTime(nb.end)}`);
          const linesToWrite = nb.lines.slice();
          if (linesToWrite.length > 0) {
            const lastIdx = linesToWrite.length - 1;
            linesToWrite[lastIdx] = (linesToWrite[lastIdx] || "").replace(/[，,]+$/,"{}");
            // 清理替换产生的空占位
            linesToWrite[lastIdx] = linesToWrite[lastIdx].replace(/\{\}/g, "");
          }
          for (const ln of linesToWrite) out.push(ln);
          out.push("");
        }
        fs.writeFileSync(srtPath, out.join("\r\n"), "utf8");
      } catch (e) {
        console.warn(`⚠️ 字幕后处理失败，已保留原始 SRT：${e.message}`);
      }
      console.log(`✅ 已生成: ${srtPath}`);
      process.exit(0);
    } else {
      console.error("❌ 未找到生成的 SRT 文件");
      process.exit(2);
    }
  });
}

main().catch((err) => {
  console.error(err);
  process.exit(1);
});


